from nx_kg_vis import app, init, NetworkXKG
from typing import Literal, get_args

NodeTypes = Literal["function", "resource", "file"]
EdgeTypes = Literal["calls", "defines_in"]


class MyKG(NetworkXKG[NodeTypes, EdgeTypes]):
    def construct(self):
        node_a_id = self.ensure_node("function", "A", file="add_node.c")
        node_b_id = self.ensure_node("function", "B", file="add_node.c")
        node_file_id = self.ensure_node("file", "add_node.c", path="add_node.c")
        self.add_edge(node_a_id, node_b_id, "calls")
        self.add_edge(node_a_id, node_file_id, "defines_in")
        self.add_edge(node_b_id, node_file_id, "defines_in")

kg = MyKG.create(list(get_args(NodeTypes)), list(get_args(EdgeTypes)))
kg.construct()
init(None, None, kg)
app.run(host="0.0.0.0", port=5000, debug=True)
